package org.hepeng.workx.spring.cloud.netflix.ribbon.loadbalancer;

import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.RoundRobinRule;
import com.netflix.loadbalancer.Server;
import com.netflix.niws.loadbalancer.DiscoveryEnabledServer;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.math.BigDecimal;
import java.util.List;
import java.util.Map;

/**
 * @author he peng
 */
public class BestMemoryAvailableRule extends RoundRobinRule {

    private BestMemoryRuleUsedAdjudicator adjudicator = key -> true;

    public BestMemoryAvailableRule() {}

    public BestMemoryAvailableRule(BestMemoryRuleUsedAdjudicator adjudicator) {
        this.adjudicator = adjudicator;
    }

    @Override
    public Server choose(ILoadBalancer lb, Object key) {

        List<Server> allServers = lb.getAllServers();
        if (CollectionUtils.isEmpty(allServers)) {
            return null;
        }

        boolean judgmentResult = this.adjudicator.adjudicate(key);
        if (! judgmentResult) {
            return super.choose(lb , key);
        }

        BigDecimal minimalMemoryUsage = BigDecimal.ONE;
        Server bestMemServer = null;
        for (Server server : allServers) {
            DiscoveryEnabledServer discoveryServer = (DiscoveryEnabledServer) server;
            Map<String, String> metadata = discoveryServer.getInstanceInfo().getMetadata();
            String memTotal = metadata.get("mem.total");
            String memUsed = metadata.get("mem.used");

            if (StringUtils.isNotBlank(memTotal) && StringUtils.isNotBlank(memUsed)) {
                BigDecimal memUsage = new BigDecimal(memUsed).divide(new BigDecimal(memTotal));

                if (memUsage.compareTo(minimalMemoryUsage) < 0) {
                    minimalMemoryUsage = memUsage;
                    bestMemServer = server;
                }
            }

        }

        return bestMemServer;
    }

    public interface BestMemoryRuleUsedAdjudicator {

        boolean adjudicate(Object key);
    }

}
